package org.xteam.sled.solver;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.xteam.sled.semantic.exp.DefaultExpVisitor;
import org.xteam.sled.semantic.exp.Exp;
import org.xteam.sled.semantic.exp.ExpAdd;
import org.xteam.sled.semantic.exp.ExpCondition;
import org.xteam.sled.semantic.exp.ExpConst;
import org.xteam.sled.semantic.exp.ExpDiv;
import org.xteam.sled.semantic.exp.ExpMod;
import org.xteam.sled.semantic.exp.ExpSub;
import org.xteam.sled.semantic.exp.ExpVar;
import org.xteam.sled.semantic.exp.TermList;
import org.xteam.sled.semantic.exp.VariableSubstitute;


public class EquationSolver {
	
	private static class VariablesCollector extends DefaultExpVisitor {
		
		Set<String> originalVariables;
		
		public VariablesCollector(Set<String> inputs) {
			originalVariables = new HashSet<String>(inputs);
		}
		
		public void visitVar(ExpVar expVar) {
			originalVariables.add(expVar.name());
		}
	}
	
	private boolean debug = false;
	private Set<String> inputs;
	
	public SolverResult solve(Set<String> inputs, List<ExpCondition> equations, boolean debug) {
		boolean d = this.debug;
		this.debug = debug;
		SolverResult res = solve(inputs, equations);
		this.debug = d;
		return res;
	}

	public SolverResult solve(Set<String> inputs, List<ExpCondition> equations) {
		debug("=== solve " + equations);
		this.inputs = inputs;
		VariablesCollector collect = new VariablesCollector(inputs);
		for (ExpCondition eqn : equations) {
			eqn.visit(collect);
		}
		debug("originalVariables=" + collect.originalVariables);
		FreshState fresh = new FreshState(collect.originalVariables);
		List<ExpCondition> constraints = new ArrayList<ExpCondition>();
		List<ExpCondition> zeros = new ArrayList<ExpCondition>();
		BalanceInfo info = new BalanceInfo(new Balancer(fresh, debug).balanceEqns(equations, constraints, zeros));
		Map<String, Exp> value = new HashMap<String, Exp>();
		for (String input : inputs) {
			value.put(input, new ExpVar(input));
			info.makeKnown(input);
		}
		List<TermList> r = new ArrayList<TermList>();
		for (ExpCondition tl : zeros) {
			r.add(tl.toTermList());
		}
		CompleteBalanceInfo cp = new CompleteBalanceInfo(value, constraints, new ArrayList<TermList>(), r);
		completeBalances(cp, info);
		return solve(r, constraints, cp.value, fresh, info, inputs);
	}
	
	private SolverResult solve(List<TermList> zeros, List<ExpCondition> constraints,
			Map<String, Exp> value, FreshState fresh, BalanceInfo info, Set<String> inputs) {
		List<TermList> pending = new ArrayList<TermList>();
		while (zeros.size() > 0 || pending.size() > 0) {
			
			if (zeros.size() > 0) {
				debug("=== state " + constraints + " " + value + " " + info.getBalances());
				TermList tl = zeros.remove(0);
				debug("=== computing zero " + tl);
				tl.normalize();
				String v = tl.candidate(inputs);
				if (v != null) { 
					Exp vpsVal = tl.eliminate(v).toSum();
					value.put(v, vpsVal);
					debug("new value " + v + " -> " + vpsVal);
					pending.addAll(zeros);
					zeros.clear();
					VariableSubstitute vsub = new VariableSubstitute(v, vpsVal);
					for (TermList r : pending) {
						zeros.add(vsub.rewrite(r.toSum()).toTermList());
					}
					pending.clear();
					for (String var : value.keySet()) {
						value.put(var, vsub.rewrite(value.get(var)));
					}
					if (new IsComputable(inputs).get(vpsVal)) {
						debug(v + " is computable");
						info.makeKnown(v);
						completeBalances(new CompleteBalanceInfo(value, constraints, pending, zeros), info);
					}
				} else {
					if (tl.hasUnknownVariable(inputs)) {
						debug("add pending " + tl);
						// z has an unknown variable
						pending.add(tl);
					} else {
						// should try to check if [[z]] is satisfiable
						constraints.add(tl.toConstraint());
					}
				}
			} else if (pending.size() > 0) {
				debug("computing pendings " + pending);
				SeparationResult res = separateNoninput(pending, inputs);
				if (res == null) {
					List<ExpCondition> r = new ArrayList<ExpCondition>();
					for (TermList tl : pending) {
						r.add(tl.toConstraint());
					}
					return new SolverResult(r);
				} else {
					// k is the coefficient of v in z, which is removed from pending
					int k = - res.factor;
					Exp z = new ExpAdd(res.termList.toSum(), new TermList.Term(k, new ExpVar(res.name)).toExpression(false));
					if (k < 0) {
						k = -k;
						z = new ExpSub(new ExpConst(0), z);
					}
					if (k <= 1)
						throw new RuntimeException("impossible coefficient");
					List<ExpCondition> equations = new ArrayList<ExpCondition>();
					equations.add(new ExpCondition(ExpCondition.EQ, new ExpVar(res.name), new ExpDiv(k, z)));
					equations.add(new ExpCondition(ExpCondition.EQ, new ExpConst(0), new ExpMod(k, z)));
					List<ExpCondition> p = new ArrayList<ExpCondition>();
					List<Balance> newbals = new Balancer(fresh, debug).balanceEqns(equations, constraints, p);
					for (ExpCondition e : p) {
						pending.add(e.toTermList());
					}
					info.getBalances().addAll(newbals);
					completeBalances(new CompleteBalanceInfo(value, constraints, new ArrayList<TermList>(), pending), info);
					zeros = pending;
					pending = new ArrayList<TermList>();
				}
			}
		}
		Iterator<String> i = value.keySet().iterator();
		while (i.hasNext()) {
			String name = i.next();
			if (! fresh.isOriginal(name) || inputs.contains(name))
				i.remove();
		}
		return new SolverResult(constraints, value);
	}
	
	private void debug(String msg) {
		if (debug)
			System.out.println(msg);
	}

	/** find a pending equation which can be further processed
	 * 
	 * @param pending
	 * @param inputs
	 * @return
	 */
	private SeparationResult separateNoninput(List<TermList> pending, Set<String> inputs) {
		for (TermList tl : pending) {
			SeparationResult res = tl.findNonInputTerm(inputs);
		 	if (res != null) {
		 		pending.remove(tl);
		 		return res;
		 	}
		}
		return null;
	}

	/*private void parseEquations(List<ExpCondition> equations, List<ExpCondition> constraints,
			List<ExpCondition> zeros) {
		for (ExpCondition equation : equations) {
			if (equation.condition() == ExpCondition.EQ) {
				zeros.add(equation); //new ExpSub(equation.left(), equation.right()).toTermList());
			} else {
				constraints.add(equation);
			}
		}
	}*/

	private void completeBalances(CompleteBalanceInfo cp, BalanceInfo info) {
		while (info.getNewlyKnown().size() > 0) {
			for (String var : new HashSet<String>(info.getNewlyKnown())) {
				maybeCompleteBalanceWithV(var, cp, info);
			}
		}
	}

	private void maybeCompleteBalanceWithV(String var, CompleteBalanceInfo cp, BalanceInfo info) {
		info.makeAlreadyKnown(var);
		completeVpsBalances (var, info, cp);
		IsComputable computable = new IsComputable(inputs);
		for (String v : cp.value.keySet()) {
			if (computable.get(cp.value.get(v))) {
				info.makeKnown(v);
			}
		}
	}
	
	private static class IsComputable extends DefaultExpVisitor {
		boolean isComputable;
		Set<String> inputs;
		public IsComputable(Set<String> inputs) {
			this.inputs = inputs;
		}
		public boolean get(Exp exp) {
			isComputable = true;
			exp.visit(this);
			return isComputable;
		}
		public void visitVar(ExpVar var) {
			isComputable &= inputs.contains(var.name());
		}
	}

	private void completeVpsBalances(String var, BalanceInfo info,
			CompleteBalanceInfo cp) {
		Iterator<Balance> i = info.getBalances().iterator();
		while (i.hasNext()) {
			Balance bal = i.next();
			if (bal.hasVariable(var)) {
				List<BalanceItem> known = null, unknown = null;
				if (allKnown(bal.getLeft(), cp.value)) {
		        	  known=bal.getLeft();
		        	  unknown=bal.getRight();
				} else if (allKnown(bal.getRight(), cp.value)) {
					known=bal.getRight();
					unknown=bal.getLeft();
				}
				if (known != null && unknown != null) {
					i.remove();
					makeUnknownsKnown(known, unknown, info, cp);
				}
			}
		}
	}

	private void makeUnknownsKnown(List<BalanceItem> known, List<BalanceItem> unknown,
			BalanceInfo info, CompleteBalanceInfo cp) {
		for (BalanceItem item : unknown) {
			debug("make known " + item);
			Exp e = new VariableSubstitute(cp.value).rewrite(item.getExpr());
			rebindValue (item.variable(), e, cp);
			if (! new IsComputable(inputs).get(e))
				throw new RuntimeException("balance var not known");
			info.makeKnown(item.variable());
		}
	}

	private void rebindValue(String variable, Exp e, CompleteBalanceInfo cp) {
		Exp ep = cp.value.get(variable);
		if (ep != null) {
			if (new IsComputable(inputs).get(ep)) {
				cp.constraints.add(new ExpCondition(ExpCondition.EQ, e, ep));
			} else {
				cp.value.put(variable, e);
				cp.zeroes.add(new ExpCondition(ExpCondition.EQ, e, ep).toTermList());
       /*    
	     val _ = let fun uFree e = Exp.isFreeIn(u, e)
		         fun olcFree olc = List.exists (fn (_, e) => uFree e) olc
		     in  if List.exists olcFree zeroes orelse List.exists olcFree pending
		         orelse StringMap.foldl (fn (e, free) => free orelse uFree e) false value
		         then
		          Impossible.impossible "u a free variable?!"
		         else
		           ()
		 */
			}
		} else {
			cp.value.put(variable, e);
			VariableSubstitute subs = new VariableSubstitute(variable, e);
			for (int i = 0; i < cp.zeroes.size(); ++i)
				cp.zeroes.set(i, subs.rewrite(cp.zeroes.get(i).toSum()).toTermList());
			for (int i = 0; i < cp.constraints.size(); ++i)
				cp.constraints.set(i, (ExpCondition)subs.rewrite(cp.constraints.get(i)));
			for (int i = 0; i < cp.pendings.size(); ++i)
				cp.pendings.set(i, subs.rewrite(cp.pendings.get(i).toSum()).toTermList());
			for (String var : cp.value.keySet()) {
				cp.value.put(var, subs.rewrite(cp.value.get(var)));
			}
/*
	     val update = Exp.subst(u, e)
	     val updateOlc = Exp.olc o update o Exp.sum
	     val zeroes = map updateOlc zeroes
	     val constraints = map (Exp.cmap update) constraints
	     val pending = map updateOlc pending
	     val value = StringMap.map update value */
	
		}
	}

	private boolean allKnown(List<BalanceItem> side, Map<String, Exp> value) {
		for (BalanceItem obj : side) {
			Exp e = value.get(obj.variable());
			if (e == null) {
				return false;
			}
			if (! new IsComputable(inputs).get(e)) {
				return false;
			}
		}
		return true;
	}

	private static class CompleteBalanceInfo {
		
		private Map<String, Exp> value;
		private List<ExpCondition> constraints;
		private List<TermList> pendings;
		private List<TermList> zeroes;
		
		public CompleteBalanceInfo(Map<String, Exp> value,
				List<ExpCondition> constraints, List<TermList> pendings, List<TermList> zeroes) {
			this.value = value;
			this.constraints = constraints;
			this.zeroes = zeroes;
			this.pendings = pendings;
		}

	}

}
